from langchain.agents import create_react_agent, AgentExecutor
from langchain_community.document_loaders import JSONLoader
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.llms.tongyi import Tongyi
from langchain_community.vectorstores import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import Tool
from langchain_text_splitters import CharacterTextSplitter
import os

def add_db():
    doc1 = JSONLoader(
        file_path="./data.json",
        jq_schema=".[].content +.[].answer",
        text_content=True,
    ).load()

    splitter = CharacterTextSplitter('\n', chunk_size=100, chunk_overlap=0)
    chunks = splitter.split_documents(doc1)
    embeddings = DashScopeEmbeddings()

    Chroma.from_documents(chunks, embedding=embeddings, persist_directory="./home/chroma_db")

def search_db(input):
    embeddings = DashScopeEmbeddings()
    db = Chroma(persist_directory="./home/chroma_db", embedding_function=embeddings)
    res = db.similarity_search(input, k=3)
    data = ""
    for i in res:
        data += i.page_content + "\n"
    return data


tools = [
    Tool(func=search_db, name="search_db", description="搜索数据库")
]
def get_message_chroma(user_input):
    llm = Tongyi()
    template = '''请尽可能回答以下问题。您可以使用以下工具:
    
               {tools}
    
               使用以下格式:
    
               Question: the input question you must answer
               Thought: you should always think about what to do
               Action: the action to take, should be one of [{tool_names}]
               Action Input: the input to the action
               Observation: the result of the action
               ... (this Thought/Action/Action Input/Observation can repeat N times)
               Thought: I now know the final answer
               Final Answer: the final answer to the original input question
    
               Begin!
    
               Question: {input}
               Thought:{agent_scratchpad}'''
    prompt = ChatPromptTemplate.from_template(template)
    agent = create_react_agent(tools=tools, llm=llm, prompt=prompt)
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
    ret = agent_executor.invoke({'input': user_input})
    return ret
